﻿/*
	Copyright (c) 2013 Andrew Meyer

	Permission is hereby granted, free of charge, to any person obtaining a copy of this software
	and associated documentation files (the "Software"), to deal in the Software without
	restriction, including without limitation the right to use, copy, modify, merge, publish,
	distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the
	Software is furnished to do so, subject to the following conditions:

	The above copyright notice and this permission notice shall be included in all copies or
	substantial portions of the Software.

	THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING
	BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
	NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
	DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
	OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
*/

using System;
using System.Collections.Generic;
using System.Collections.Specialized;
using System.Linq;
using System.Threading.Tasks;
using System.Windows;
using System.Windows.Media.Media3D;
using System.Windows.Threading;

namespace Ascend
{
	/// <summary>
	/// Class that describes the bone structure for an Ascend.AnimatedSceneNode. This class cannot
	/// be inherited.
	/// </summary>
	public sealed class Armature : DependencyObject
	{
		#region Lifecycle
		/// <summary>
		/// Initializes a new instance of the Ascend.Armature class.
		/// </summary>
		public Armature()
		{
			Data = new Model3DGroup();

			// Create root Bone model group
			Data.Children.Add(new Model3DGroup());

			// Create root Bone collection. Each element will correspond to a model in the root Bone
			// model group.
			RootBones = new BoneCollection();

			// Listen to changes in root Bones
			RootBones.CollectionChanged += RootBones_CollectionChanged;
		}
		#endregion

		#region Private Members
		internal Model3DGroup Data
		{
			get
			{
				return (Model3DGroup)GetValue(DataProperty);
			}

			set
			{
				SetValue(DataProperty, value);
			}
		}

		private static readonly DependencyProperty DataProperty =
			DependencyProperty.Register("Data", typeof(Model3DGroup), typeof(Armature));


		private Model3DCollection VisualChildren
		{
			get
			{
				return ((Model3DGroup)Data.Children[0]).Children;
			}
		}
		#endregion

		#region Public Properties
		/// <summary>
		/// Gets or sets the name of this Ascend.Armature.
		/// </summary>
		public string Name
		{
			get;
			set;
		}

		internal AnimatedSceneNode _skin;
		/// <summary>
		/// Gets the Ascend.AnimatedSceneNode acting as this Ascend.Armature's skin.
		/// </summary>
		public AnimatedSceneNode Skin
		{
			get
			{
				return _skin;
			}
		}

		/// <summary>
		/// Gets this Ascend.Armature's collection of top-level Ascend.Bone instances.
		/// </summary>
		public BoneCollection RootBones
		{
			get;
			private set;
		}
		#endregion

		#region Public Methods
		/// <summary>
		/// Places each Ascend.Bones belonging to this Ascend.Armature into its rest pose and
		/// updates this Ascend.Armature's skin.
		/// </summary>
		public void Rest()
		{
			foreach (Bone bone in ListBones())
			{
				bone.Rest();
			}

			UpdateSkin();
		}

		/// <summary>
		/// Updates this Ascend.Armature's skin for its current bone poses.
		/// </summary>
		public void UpdateSkin()
		{
			UpdateSkinForPoses();
		}

		/// <summary>
		/// Creates a deep copy of this Ascend.Armature, including deep copies of all Ascend.Bone
		/// instances.
		/// </summary>
		/// <returns>
		/// Returns the deep copy.
		/// </returns>
		/// <remarks>
		/// Since an Ascend.AnimatedSceneNode may only have one Ascend.Armature associated with it,
		/// the clone's Skin property will be null.
		/// </remarks>
		public Armature Clone()
		{
			Armature clone = new Armature
			{
				Name = Name
			};

			foreach (Bone bone in RootBones)
			{
				clone.RootBones.Add((Bone)bone.Clone());
			}

			return clone;
		}

		/// <summary>
		/// Creates a string representation of this Ascend.Armature.
		/// </summary>
		/// <returns>
		/// Returns a string containing this Ascend.Armature's name, if any, and type.
		/// </returns>
		public override string ToString()
		{
			return (string.IsNullOrWhiteSpace(Name)
				? string.Empty
				: Name + " ")
				+ "(" + GetType().Name + ")";
		}
		#endregion

		#region Internal Methods
		internal List<Bone> ListBones()
		{
			List<Bone> bones = new List<Bone>();

			foreach (Bone bone in RootBones)
			{
				ListBonesRecursive(bone, bones);
			}

			return bones;
		}

		internal void ApplyBonePoses(Dictionary<Bone, ObjectPose> bonePoses)
		{
			// Update skin before updating bones so that rigid attachments do not lag
			UpdateSkinForPoses(bonePoses);

			Dispatcher.BeginInvoke((Action)(() =>
				{
					foreach (Bone bone in bonePoses.Keys)
					{
						if (bonePoses.ContainsKey(bone))
						{
							ObjectPose pose = bonePoses[bone];

							if (pose != null)
							{
								bone.Position = new Point3D(pose.X, pose.Y, pose.Z);
								bone.Rotation = new Angle3D(pose.AngleX, pose.AngleY, pose.AngleZ);
								bone.Scale = new Scale3D(pose.ScaleX, pose.ScaleY, pose.ScaleZ);
							}
						}
					}
				}));
		}

		internal void ClearSkin()
		{
			if (Skin != null)
			{
				Skin.RemoveArmature();
			}
		}
		#endregion

		#region Private Methods
		private void UpdateSkinForPoses(Dictionary<Bone, ObjectPose> bonePoses = null)
		{
			if (!CheckCanUpdateSkin())
			{
				return;
			}
			
			Dictionary<int, DeltaTransform> vertexTransforms = GetVertexTransforms(bonePoses);
			MeshGeometry3D mesh = GetMesh();
			List<Point3D> positions = GetPositions(mesh);
			List<Vector3D> normals = GetNormals(mesh);

			TransformVertices(positions, normals, vertexTransforms);
			
			Dispatcher.Invoke(() =>
				{
					Skin.Model.Geometry = null;
					mesh.Positions = new Point3DCollection(positions);
					mesh.Normals = new Vector3DCollection(normals);
					Skin.Model.Geometry = mesh;
				});
		}
		
		private bool CheckCanUpdateSkin()
		{
			bool canUpdateSkin = false;

			Dispatcher.Invoke(() =>
				{
					canUpdateSkin = (Skin != null && Skin.Model != null &&
						Skin.Model.Geometry != null);
				});

			return canUpdateSkin;
		}

		private Dictionary<int, DeltaTransform> GetVertexTransforms(
			Dictionary<Bone, ObjectPose> bonePoses)
		{
			Dictionary<int, DeltaTransform> vertexTransforms =
				new Dictionary<int, DeltaTransform>();

			Dictionary<Bone, HomogeneousTransform> restTransforms =
				new Dictionary<Bone, HomogeneousTransform>();

			Dictionary<Bone, HomogeneousTransform> currTransforms =
				new Dictionary<Bone, HomogeneousTransform>();

			ComputeVertexTransforms(RootBones, vertexTransforms, restTransforms, currTransforms,
				bonePoses);

			return vertexTransforms;
		}

		private void ComputeVertexTransforms(BoneCollection bones,
			Dictionary<int, DeltaTransform> vertexTransforms,
			Dictionary<Bone, HomogeneousTransform> restTransforms,
			Dictionary<Bone, HomogeneousTransform> currTransforms,
			Dictionary<Bone, ObjectPose> bonePoses)
		{
			foreach (Bone bone in bones)
			{
				HomogeneousTransform restTrans = null;
				HomogeneousTransform currTrans = null;

				HomogeneousTransform storedTrans = (bonePoses == null)
							? null
							: new HomogeneousTransform(bonePoses[bone]);

				Dispatcher.Invoke(() =>
					{
						// If bonePoses is null, use current pose. Otherwise, use provided pose.
						currTrans = (bonePoses == null)
							? new HomogeneousTransform(bone.Data.Transform.Value)
							: storedTrans;

						restTrans = new HomogeneousTransform(bone._restTransform.Value);
					});

				if (bone.Parent != null)
				{
					currTrans = currTransforms[bone.Parent] * currTrans;
					restTrans = restTransforms[bone.Parent] * restTrans;
				}

				currTransforms[bone] = currTrans;
				restTransforms[bone] = restTrans;

				ComputeBoneVertexTransform(bone, vertexTransforms, restTransforms, currTransforms);
				ComputeVertexTransforms(bone.Children, vertexTransforms, restTransforms,
					currTransforms, bonePoses);
			}
		}

		private void ComputeBoneVertexTransform(Bone bone,
			Dictionary<int, DeltaTransform> vertexTransforms,
			Dictionary<Bone, HomogeneousTransform> restTransforms,
			Dictionary<Bone, HomogeneousTransform> currTransforms)
		{
			HomogeneousTransform vertexTrans = currTransforms[bone] * restTransforms[bone].Inverse;

			foreach (VertexInfluence inf in bone.VertexInfluences.Where(i => i.Weight != 0.0))
			{
				foreach (int vert in inf.Vertices)
				{
					if (vertexTransforms.ContainsKey(vert))
					{
						vertexTransforms[vert].Add(inf.Weight, vertexTrans);
					}
					else
					{
						vertexTransforms.Add(vert, new DeltaTransform(inf.Weight, vertexTrans));
					}
				}
			}
		}

		private MeshGeometry3D GetMesh()
		{
			MeshGeometry3D mesh = null;

			Dispatcher.Invoke(() =>
				{
					mesh = (MeshGeometry3D)Skin.Model.Geometry;
				});

			return mesh;
		}

		private List<Point3D> GetPositions(MeshGeometry3D mesh)
		{
			List<Point3D> positions = null;

			Dispatcher.Invoke(() =>
				{
					positions = new List<Point3D>(mesh.Positions);
				});

			return positions;
		}

		private List<Vector3D> GetNormals(MeshGeometry3D mesh)
		{
			List<Vector3D> normals = null;

			Dispatcher.Invoke(() =>
				{
					normals = new List<Vector3D>(mesh.Normals);
				});

			return normals;
		}

		private void TransformVertices(List<Point3D> positions, List<Vector3D> normals,
			Dictionary<int, DeltaTransform> vertexTransforms)
		{
			Parallel.ForEach(vertexTransforms.Keys, (vert) =>
			{
				bool hasNormal = (Skin._originalNormals.Count > vert);
				Vector3D inNormal = hasNormal
					? Skin._originalNormals[vert]
					: default(Vector3D);

				if (vertexTransforms.ContainsKey(vert))
				{
					HomogeneousTransform vertexTrans = vertexTransforms[vert].Value;

					positions[vert] = vertexTrans * Skin._originalVertexPositions[vert];

					if (hasNormal)
					{
						normals[vert] = vertexTrans * Skin._originalNormals[vert];
					}
				}
			});
		}

		private void ListBonesRecursive(Bone bone, List<Bone> bones)
		{
			bones.Add(bone);

			foreach (Bone child in bone.Children)
			{
				ListBonesRecursive(child, bones);
			}
		}
		#endregion

		#region Event Handlers
		private void RootBones_CollectionChanged(object sender, NotifyCollectionChangedEventArgs e)
		{
			switch (e.Action)
			{
				case NotifyCollectionChangedAction.Add:
					HandleCollectionActionAdd(sender, e);
					break;

				case NotifyCollectionChangedAction.Remove:
					HandleCollectionActionRemove(sender, e);
					break;
			}
		}

		private void HandleCollectionActionAdd(object sender, NotifyCollectionChangedEventArgs e)
		{
			foreach (Bone bone in e.NewItems)
			{
				// Clear old parent
				bone.ClearTransformParent();

				// Add data to model group
				VisualChildren.Add(bone.Data);

				// Set armature and transform parent (even if null)
				bone._armature = this;
				bone._transformParent = Skin;
			}
		}

		private void HandleCollectionActionRemove(object sender, NotifyCollectionChangedEventArgs e)
		{
			foreach (Bone bone in e.OldItems)
			{
				// Remove data from model group
				VisualChildren.Remove(bone.Data);

				// Clear armature and transform parent
				bone._armature = null;
				bone._transformParent = null;
			}
		}
		#endregion
	}
}
